import struct,time,zlib,logging,os
import rsa, fast
from binascii import b2a_hex
try: import cStringIO as strio
except ImportError: import StringIO as strio
__author__ = "lyricconch@gmail.com"
__version__ = "1.0b"


def timeit(fn):
    def dummy(*args,**kwargs):
        t= time.clock()
        ret= fn(*args,**kwargs)
        print fn.func_name, time.clock()-t  
        return ret
    return dummy  

#@timeit
def packit(buffer, package, key="", session=""):
    shead= strio.StringIO(); sbody= strio.StringIO();
    for entity in package:
        head,body= "\x00\x00".join(entity[0:3]), entity[3]
        shead.write(struct.pack(">II",len(body),len(head)))
        shead.write(head)      
        sbody.write(body)
    head= zlib.compress(shead.getvalue()) 
    body= sbody.getvalue()  
    session= "" if not key else os.urandom(32) if session is None else session  
    content= enaes(head+body, session)
    prefix= struct.pack(">HH",len(package),len(head))
    chrome= enrsa(prefix+session,key)    
    buffer.write(chrome+content)  
    #logging.info("packit: prefix=%s, deadline=%.3f"%(
    #    b2a_hex(prefix), 2+1.384**(5-len(package)) ))
    return session

#@timeit
def unpack(buffer, package, key="", session=""):
    chrome=dersa(buffer.read(64),key); rebuild=False;
    pkglen, offset= struct.unpack(">HH", chrome[:4])
    if len(package)!=pkglen: rebuild=True; del package[:]; 
    session= "" if not key else chrome[4:] if session is None else session  
    content= deaes(buffer.read(),session)
    shead= strio.StringIO(zlib.decompress(content[:offset]))
    sbody= strio.StringIO(content[offset:])   
    for index in xrange(pkglen):
        blen,hlen= struct.unpack(">II", shead.read(8))
        head= shead.read(hlen).split("\x00\x00",2)
        body= sbody.read(blen)
        if rebuild: head.append(body); package.append(head); 
        else: z=package[index]; z[:3]=head;z[3]=body;z[4].set();    
    #logging.info("unpack: prefix=%s, timeline=%.3f"%(
    #    b2a_hex(chrome[:4]), time.clock() ))
    return session
    
def str2hdr(head, MOD={}):
    if isinstance(head, dict): 
        head= hdr2str(head)
    headers= {}
    for line in head.split("\n"):
        split= line.split(":", 1)
        if len(split)!=2: continue;
        hname= split[0].lower()
        headers[hname]= split[1].strip()     
    for h,v in MOD.iteritems():
        if v is not None: headers[h]=v;
        elif h in headers: del headers[h];
    return headers        

def hdr2str(headers, MOD={}):  
    if isinstance(headers, str): 
        headers= str2hdr(headers) 
    head= []
    for h,v in MOD.iteritems():
        if v is not None: headers[h]=v;
        elif h in headers: del headers[h];
    for hname, hvalue in headers.iteritems():
        head.append(hname.title())
        head.append(": ")
        head.append(str(hvalue).strip())
        head.append("\r\n")       
    return "".join(head)

def enrsa(data, rsakey):
    if not rsakey: return data;
    N, E= rsakey[0], rsakey[1] or rsakey[2]
    return rsa.translate(data, N, E, False)[0]

def dersa(data, rsakey):
    if not rsakey: return data;
    N, D= rsakey[0], rsakey[2] or rsakey[1]    
    return rsa.translate(data, N, D, True)[0]

from Crypto.Cipher import AES
def enaes(data, aeskey):
    if not aeskey or not data: return data;
    pad= -(len(data)+1)%16
    data= chr(pad)+os.urandom(pad)+data
    return AES.new(aeskey).encrypt(data)

def deaes(data, aeskey):
    if not aeskey or not data: return data;
    data= AES.new(aeskey).decrypt(data)
    pad= ord(data[0])+1
    return data[pad:]
    

def cfast(data, fastkey):
    if not fastkey: return data;
    return fast.transBox(data, fastkey)

def newkey(len):
    return fast.randomkey(len, fast.KEY_BOX)

def invert(key):
    return fast.invertkey(key, fast.KEY_BOX)